#include "misc.h"
#include "matrix.h"
#include "util.h"

#define MAX(a,b)  (((a) > (b)) ? (a) : (b))
#define TOLERANCE   1.0e-6



/* Reads a file into a matrix */
void read_file(matrix *A, const char *filename) {
    FILE * fp = fopen(filename, "r");
    load_mat(fp, A);
    fclose(fp);
}

void calc_Q(matrix *M, matrix *HpInv, matrix*Q) {
    matrix temp, Mt;
    trans_mat(M, &Mt);

    mult_mat(M, HpInv, &temp);
    mult_mat(&temp, &Mt, Q);

    dest_mat(&temp);
    dest_mat(&Mt);
}

void calc_wXX(matrix *A, matrix *x0, matrix *xXX, matrix *d, int N, int pnum, double uXX, matrix *wXX) {

    matrix _A[N];     // _A[i] = A^i (A^0 is redundant, but useful in loops)

    /* We need A^i multiple times. So, let's pre-calculate it once */
    init_mat(&(_A[0]),A->row,A->col);  // Take the size from matrix A
    fill_ident(&(_A[0]));
    for (int i=1; i<N; i++) {
        mult_mat(&(_A[i-1]),A,&(_A[i]));
    }

    init_mat(wXX,4*(N-1),1);

    for (int i=1; i<=N-1; i++) {

        /* Calculating (xXX - A^i*x0 - sum(var=1..i)[A^(i-var)*d(var-1)] */
        /* d(var-1) is for first problem. It shifts by one for every problem */
        matrix temp1, temp2, temp3, temp4, temp5;

        /* 2nd term */
        mult_mat(&(_A[i]),x0,&temp1);

        /* 3rd term */
        init_mat(&temp2,3,1);
        fill_zero(&temp2);

        for (int var=1;var<=i;var++) {

            /* set up d(var) */
            matrix tempdvar;
            init_mat(&tempdvar,3,1); 
            fill_zero(&tempdvar);
            MAT_ELEM(&tempdvar,0,0) = MAT_ELEM(d,((pnum+var-1)/d->col),((pnum+var-1)%d->col));            

            matrix tempmult;
            mult_mat(&(_A[i-var]),&tempdvar,&tempmult);
                   
            matrix temp__;
            copy_mat(&temp2,&temp__);
            dest_mat(&temp2); 
            add_mat(&temp__,&tempmult,&temp2);

            /* cleanup*/
            dest_mat(&tempdvar);
            dest_mat(&tempmult);
            dest_mat(&temp__);
        }

        /* temp1 is 2nd term, temp2 is 3rd term */
        add_mat(&temp1,&temp2,&temp3);

        mult_scalar(&temp3,-1.0,&temp4);
        add_mat(xXX, &temp4, &temp5);        

        for (int j=0; j<3; j++)
            MAT_ELEM(wXX,(3*(i-1)+j),0) = MAT_ELEM(&temp5,j,0);
        MAT_ELEM(wXX,3*(N-1)+(i-1),0) = uXX;

        /* Cleanup */
        dest_mat(&temp1);
        dest_mat(&temp2);
        dest_mat(&temp3);
        dest_mat(&temp4);
        dest_mat(&temp5);
    }
 
    for (int i=0; i<N; i++)
        dest_mat(&(_A[i]));

}

void calc_b(matrix *A, matrix *x0, matrix *xlo, matrix *xhi, matrix *d, int N, int pnum, double ulo, double uhi, matrix *b) {

    matrix wlo, whi;
    /* calculate wlo*/
    calc_wXX(A, x0, xlo, d, N, pnum, ulo, &wlo);
    /* calculate whi*/
    calc_wXX(A, x0, xhi, d, N, pnum, uhi, &whi);

    matrix temp;
    mult_scalar(&wlo,-1.0,&temp);
   
    concat_mat_ver(&whi,&temp,b); 

    dest_mat(&wlo);
    dest_mat(&whi);
    dest_mat(&temp);
}

/* Equation 3 from the paper */
void calc_eq3(matrix *lambda, matrix *Q, matrix *h, double *F) {

    matrix temp1, temp2, temp3, temp4;
    matrix Fmat;
    matrix lambdat;
  
    /* First term */ 
    trans_mat(lambda, &lambdat);
    mult_mat(&lambdat, Q, &temp1);
    mult_mat(&temp1, lambda, &temp2);
    mult_scalar(&temp2, 0.5, &temp3);

    /* Second term */
    mult_mat(&lambdat, h, &temp4);

    /* add */
    add_mat(&temp3, &temp4, &Fmat);

    /* Now this is really a single value */
    assert(Fmat.row == 1 && Fmat.col == 1);
    *F = MAT_ELEM(&Fmat,0,0);

    dest_mat(&temp1);
    dest_mat(&temp2);
    dest_mat(&temp3);
    dest_mat(&temp4);
    dest_mat(&lambdat);
    dest_mat(&Fmat);
}

/* Equation 3 from the paper */
void calc_eq3_fast(matrix *_lambda, matrix *_Q, matrix *_h, double *F) {

    /* we do stuff the dirty way, for performance */

    int n = _lambda->row;

    double *y = _lambda->data;
    double *Q = _Q->data;
    double *h = _h->data;
    assert(_Q->major == ROWMAJOR);

    double *temp = (double *)calloc(n,sizeof(double));

    /* Calculate Q*y first, temp = Q*y*/
    for (int i=0; i<n; i++) {
        /* temp[i] calculation */
        for (int j=0; j<n; j++) {
            temp[i] += Q[(i*n)+j]*y[j];
        }
    }

    double t1 = 0.0;
    double t2 = 0.0;
    /* calculate y^T * temp */
    for (int i=0; i<n; i++)
        t1 += y[i]*temp[i];

    for (int i=0; i<n; i++)
        t2 += y[i]*h[i];

    *F = 0.5*t1 + t2;

    free(temp);
}



/* This calculates r to be used for the calculation of Q- and Q+ */
void calc_r(matrix *Q, matrix *r) {
/*
    matrix temp1, temp2; 
    mult_scalar(Q, -1.0, &temp1);

    // Update temp1 to max(-Q and zero) 
    for (int i=0; i<(temp1.row); i++) 
        for (int j=0; j<(temp1.col); j++) 
            MAT_ELEM(&temp1,i,j) = MAX(MAT_ELEM(&temp1,i,j),0.0);

    init_mat(&temp2,Q->row,1);
    for (int i=0;i<(temp2.row);i++)
        MAT_ELEM(&temp2,i,0) = 1.0;        

    mult_mat(&temp1,&temp2,r);

    dest_mat(&temp1);
    dest_mat(&temp2);
*/

    init_mat(r,Q->row,1);
    fill_zero(r);

}

/* Calculate Qxx for updating lambda 
    if hi == 1, then calculates Qhi
    if hi == 0, then calculates Qlo
*/
void calc_QXX(matrix *Q, matrix *r, int hi, matrix *QXX) {

    matrix diag_r;
    init_mat(&diag_r,r->row,r->row);
    fill_zero(&diag_r);

    init_mat(QXX,Q->row,Q->col);

    for (int i=0;i<(r->row);i++)
        MAT_ELEM(&diag_r,i,i) = MAT_ELEM(r,i,0); 

    for (int i=0;i<(Q->row);i++)
        for (int j=0; j<(Q->col); j++) {

            if (hi == 1)
                MAT_ELEM(QXX,i,j) = MAX(0.0,MAT_ELEM(Q,i,j)) + MAT_ELEM(&diag_r,i,j);
            else
                MAT_ELEM(QXX,i,j) = MAX(0.0,-1.0*MAT_ELEM(Q,i,j)) + MAT_ELEM(&diag_r,i,j);

        }
    
    dest_mat(&diag_r);
}


/* update lambda as per eq 8 */
void update_lambda(matrix *h, matrix *Qlo, matrix *Qhi, matrix *lambda) {
    int row = lambda->row;
    int col = lambda->col;

    matrix temp1, temp2;

    mult_mat(Qlo,lambda,&temp1);
    mult_mat(Qhi,lambda,&temp2);

    for (int i=0;i<row;i++) {
        for (int j=0;j<col;j++) {

            double t1 = MAX(0.0,-1*MAT_ELEM(h,i,j)) + MAT_ELEM(&temp1,i,j);
            double t2 = MAX(0.0,MAT_ELEM(h,i,j)) + MAT_ELEM(&temp2,i,j);
            if (fabs(t2)>0.000001)
                MAT_ELEM(lambda,i,j) = MAT_ELEM(lambda,i,j)*(t1/t2);
            else
                MAT_ELEM(lambda,i,j) = 0.0;

        }
    }

    dest_mat(&temp1);
    dest_mat(&temp2);

}

/* update lambda as per eq 8 */
void update_lambda_fast(matrix *h, matrix *Qlo, matrix *Qhi, matrix *lambda) {

    int row = lambda->row;
/*
    assert(lambda->col == 1);
    assert(h->col == 1);
    assert(Qlo->major == ROWMAJOR);
    assert(Qhi->major == ROWMAJOR);
*/
    double *_Qlo = Qlo->data;
    double *_Qhi = Qhi->data;
    double *_h = h->data;
    double *y = (double *)malloc(sizeof(double)*row);     /* a copy of the data */
    memcpy(y,lambda->data,sizeof(double)*row);
    double *_y = lambda->data;  /* This is what we will update */

    // Process lambda[i]
    for (int i=0;i<row;i++) {

        if (fabs(y[i])<1.0e-10) {
            _y[i] = 0.0;            
            continue;
        }


        double t1 = 0.0;    // numerator
        double t2 = 0.0;    // denominator

        for (int j=0; j<row; j++) { 
            t1 += _Qlo[(i*row)+j]*y[j];
            t2 += _Qhi[(i*row)+j]*y[j];
        }

        t1 += MAX(0.0,-1*_h[i]);
        t2 += MAX(0.0,1*_h[i]);

        if (fabs(t2)>0.000001) 
            _y[i] = _y[i]*(t1/t2);
        else
            _y[i] = 0.0;

    }

    free(y);
}

/* calculate equation 7 from the paper */
void calc_U(matrix *HpInv, matrix *M, matrix *lambda, matrix *U) {

    matrix temp1, temp2;
    matrix Mt;

    /* The term inside the parenthesis */
    trans_mat(M,&Mt);
    mult_mat(&Mt,lambda,&temp1);

    mult_mat(HpInv, &temp1, &temp2);

    mult_scalar(&temp2, -1.0, U);

    dest_mat(&temp1);
    dest_mat(&temp2);
    dest_mat(&Mt);
    
}


/* x1 = A*x0+B*u(0) +d(0) */
void calc_x1(matrix *A, matrix *x0, matrix *B, matrix *U, matrix *d, int pnum, matrix *x1) {

    matrix temp1, temp2, temp3, temp4;

    /* First term */
    mult_mat(A,x0,&temp1);

    /* second term */
    mult_scalar(B,MAT_ELEM(U,0,0),&temp2);

    /* third term */
    init_mat(&temp3,3,1);
    fill_zero(&temp3);
    MAT_ELEM(&temp3,0,0) = MAT_ELEM(d,(pnum/d->col),(pnum%d->col));

    add_mat(&temp1,&temp2,&temp4);
    add_mat(&temp4,&temp3,x1);

    dest_mat(&temp1);
    dest_mat(&temp2);
    dest_mat(&temp3);
    dest_mat(&temp4);

}

int main (int argc, char *argv[]) {

    int n=1, N=180;
    matrix M, HpInv, Q;     // for Q
    matrix A, B, d;
    double ulo=-10.0, uhi=10.0;
    matrix xlo, xhi;

    /* Load matrices from files */
    read_file(&M, "data/M.bin");
    read_file(&HpInv, "data/HpInv.bin");    
    read_file(&A, "data/A.bin");
    read_file(&B, "data/B.bin");
    read_file(&d, "data/d.bin");    

    calc_Q(&M,&HpInv,&Q);

    /* Initialize other matrices */
    init_mat(&xlo,3,1);
    for (int i=0; i<3; i++) MAT_ELEM(&xlo,i,0) = 20.0;
    init_mat(&xhi,3,1);
    for (int i=0; i<3; i++) MAT_ELEM(&xhi,i,0) = 23.0;

    matrix x0;          // This is what links different problems together    
 
    init_mat(&x0,3,1);
    MAT_ELEM(&x0,0,0) = 22.9;
    MAT_ELEM(&x0,1,0) = 22.8;
    MAT_ELEM(&x0,2,0) = 22.8; 

    matrix Qhi, Qlo;
    matrix r;

    calc_r(&Q,&r);
    calc_QXX(&Q,&r,1,&Qhi);
    calc_QXX(&Q,&r,0,&Qlo);


    for (int pnum=0; pnum<100; pnum++) {

        double u0, _x0;
    
        matrix x1;   
        matrix b, lambda, U;
       
        calc_b(&A, &x0, &xlo, &xhi, &d, N, pnum, ulo, uhi, &b);
    
        /* Now since we have everything we need to solve the dual program, let's solve it */
        init_mat(&lambda,8*n*(N-1),1);
        for (int i=0; i<8*n*(N-1);i++)  
            MAT_ELEM(&lambda,i,0) = 1.0;   // Initialize lambda with 1's
    
        int iter=0;
        double F=0.0, Flast=0.0;
    
        double s = read_timer();
        do {
            iter++; Flast = F;
    
           (calc_eq3_fast(&lambda, &Q, &b, &F));
           (update_lambda_fast(&b, &Qlo, &Qhi, &lambda));
            //printf("(%d): Iter %d: F(lambda) = %f\n", pnum, iter, F);   
 
        } while ((iter<2) || (fabs(Flast-F)>TOLERANCE));
        double e = read_timer();
    
        // printf("F:%lf, Flast:%lf\n", F, Flast);
        
        //printf("Computation time: %f",e-s);

        calc_U(&HpInv,&M,&lambda,&U);
        u0 = MAT_ELEM(&U,0,0);    

        calc_x1(&A,&x0,&B,&U,&d,pnum,&x1);
        _x0 = MAT_ELEM(&x1,0,0);

        // printf(">>> %d: %f <<<\n", pnum, MAT_ELEM(&x0,0,0)); 

        /* Copy x1 to x0 */
        for (int i=0; i<x0.row; i++)
            MAT_ELEM(&x0,i,0) = MAT_ELEM(&x1,i,0);
 
        dest_mat(&lambda);
        dest_mat(&x1);
        dest_mat(&b);
        dest_mat(&U);


        printf("Time step %d: time: %lf, cost: %lf, iterations: %d, u(0): %lf, x0: %lf\n", pnum, e-s, F, iter, u0, _x0);

    }
    
    /* Cleanup */
    dest_mat(&Qhi);
    dest_mat(&Qlo);
    dest_mat(&r);
    dest_mat(&x0);
    dest_mat(&M);
    dest_mat(&HpInv);
    dest_mat(&Q);
    dest_mat(&xlo);
    dest_mat(&xhi);
    dest_mat(&A);
    dest_mat(&d);
    dest_mat(&B);

    return 0;
}

/* 
 lambda's size is 8*n*(N-1)
*/
